/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2024. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef GE_TILINGDATA_BASE_H
#define GE_TILINGDATA_BASE_H

#include <vector>
#include <map>
#include <memory>
#include <cstring>
#include <securec.h>
#include "graph/ascend_string.h"
#include "register/definition.h"

namespace ge {
struct CharPtrCmp {
    bool operator()(const char *strLeft, const char *strRight) const
    {
        return strcmp(strLeft, strRight) < 0;
    }
};

class GRAPH_API_EXPORT StructSizeInfoBase {
public:
    static StructSizeInfoBase &GetInstance();
    void SetStructSize(const char *structType, const size_t structSize)
    {
        if (structSizeInfo.find(structType) != structSizeInfo.end()) {
            return;
        }
        structSizeInfo[structType] = structSize;
    }
    size_t GetStructSize(const char *structType)
    {
        if (structSizeInfo.find(structType) == structSizeInfo.end()) {
            return 0;
        }
        return structSizeInfo.at(structType);
    }

private:
    StructSizeInfoBase(){};
    ~StructSizeInfoBase(){};
    StructSizeInfoBase(const StructSizeInfoBase &);
    StructSizeInfoBase &operator=(const StructSizeInfoBase &);
    std::map<const char *, size_t, CharPtrCmp> structSizeInfo;
};

class FieldInfo {
public:
    FieldInfo(const char *dtype, const char *name) : dtype_(dtype), name_(name), classType_("0")
    {}
    FieldInfo(const char *dtype, const char *name, size_t arrSize)
        : dtype_(dtype), name_(name), arrSize_(arrSize), classType_("1")
    {}
    FieldInfo(const char *dtype, const char *name, const char *structType, size_t structSize)
        : dtype_(dtype), name_(name), structType_(structType), structSize_(structSize), classType_("2")
    {}

public:
    const char *dtype_;
    const char *name_;
    size_t arrSize_;
    const char *structType_;
    size_t structSize_;
    const char *classType_;
};

class GRAPH_API_EXPORT TilingDef {
public:
    ~TilingDef()
    {
        if (!dataPtrInited_ && dataPtr_ != nullptr) {
            delete[] dataPtr_;
        }
        dataPtr_ = nullptr;
        className_ = nullptr;
    }
    void SaveToBuffer(void *pdata, size_t capacity);
    std::vector<FieldInfo> GetFieldInfo() const;
    const char *GetTilingClassName() const;
    size_t GetDataSize() const;
    void SetDataPtr(void *dataPtr);
    void CheckAlignAndGenPlaceHolder(const char *name, size_t typeSize);

protected:
    void InitData();
    void GeLogError(const std::string &str) const;
    // dtype, name
    std::vector<FieldInfo> fieldInfo_;
    uint8_t *dataPtr_ = nullptr;
    size_t dataSize_ = 0;
    const char *className_;
    std::vector<std::pair<void *, size_t>> saveBufferPtr;
    size_t structSize_ = 0;
    bool dataPtrInited_ = false;
    uint32_t featureBitFlat_ = 0;
    uint8_t reserved_buf[128] = {0};
};

using TilingDataConstructor = std::shared_ptr<TilingDef> (*)();

class GRAPH_API_EXPORT CTilingDataClassFactory {
public:
    static CTilingDataClassFactory &GetInstance();
    void RegisterTilingData(const char *opType, const char *className, const TilingDataConstructor constructor);
    std::shared_ptr<TilingDef> CreateTilingDataInstance(const char *opType);
    std::string GetTilingDataClassName(const char *opType);

private:
    CTilingDataClassFactory(){};
    ~CTilingDataClassFactory(){};
    CTilingDataClassFactory(const CTilingDataClassFactory &);
    CTilingDataClassFactory &operator=(const CTilingDataClassFactory &);
    std::map<std::string, std::pair<std::string, TilingDataConstructor>> instance_;
};
}  // end of namespace optiling

/*
example:
// supported data_type: int8_t/uint8_t/int16_t/uint16_t/int32_t/uint32_t/int64_t/uint64_t
BEGIN_TILING_DATA_DEF(MaxPoolTilingData)
    // format: TILING_DATA_FIELD_DEF(data_type, field_name);
    TILING_DATA_FIELD_DEF(int32_t, dim_0);
    TILING_DATA_FIELD_DEF(uint8_t, var_1);
    TILING_DATA_FIELD_DEF(int64_t, factor_1);
END_TILING_DATA_DEF
REGISTER_TILING_DATA_CLASS(MaxPool, MaxPoolTilingData)
*/
#define BEGIN_TILING_DATA_DEF(className)                                                                               \
    class className : public TilingDef {                                                                               \
    public:                                                                                                            \
        size_t FieldHandler(const char *dtype, const char *name, size_t typeSize, const char *namePh)                  \
        {                                                                                                              \
            CheckAlignAndGenPlaceHolder(namePh, typeSize);                                                             \
            fieldInfo_.emplace_back(FieldInfo(dtype, name));                                                           \
            size_t ret_val = dataSize_;                                                                                \
            dataSize_ += typeSize;                                                                                     \
            return ret_val;                                                                                            \
        }                                                                                                              \
        size_t FieldHandler(const char *dtype, const char *name, size_t typeSize, size_t arrSize, const char *namePh)  \
        {                                                                                                              \
            CheckAlignAndGenPlaceHolder(namePh, typeSize);                                                             \
            fieldInfo_.emplace_back(FieldInfo(dtype, name, arrSize));                                                  \
            size_t ret_val = dataSize_;                                                                                \
            dataSize_ += typeSize * arrSize;                                                                           \
            return ret_val;                                                                                            \
        }                                                                                                              \
        size_t FieldHandler(const char *dtype, const char *name, const char *structType, size_t structSize, void *ptr, \
            const char *namePh)                                                                                        \
        {                                                                                                              \
            CheckAlignAndGenPlaceHolder(namePh, 8);                                                                    \
            fieldInfo_.emplace_back(FieldInfo(dtype, name, structType, structSize));                                   \
            size_t ret_val = dataSize_;                                                                                \
            dataSize_ += structSize;                                                                                   \
            saveBufferPtr.emplace_back(std::make_pair(ptr, ret_val));                                                  \
            structSize_ += structSize;                                                                                 \
            return ret_val;                                                                                            \
        }                                                                                                              \
                                                                                                                       \
    public:                                                                                                            \
        className()                                                                                                    \
        {                                                                                                              \
            className_ = #className;                                                                                   \
            CheckAlignAndGenPlaceHolder(#className "PH", 8);                                                           \
            StructSizeInfoBase::GetInstance().SetStructSize(#className, dataSize_);                                    \
            InitData();                                                                                                \
        }                                                                                                              \
        explicit className(void *ptr)                                                                                  \
        {                                                                                                              \
            className_ = #className;                                                                                   \
            CheckAlignAndGenPlaceHolder(#className "PH", 8);                                                           \
            StructSizeInfoBase::GetInstance().SetStructSize(#className, dataSize_);                                    \
            if (ptr == nullptr) {                                                                                      \
                return;                                                                                                \
            }                                                                                                          \
            SetDataPtr(ptr);                                                                                           \
        }

#define TILING_DATA_FIELD_DEF(data_type, field_name)                                                          \
public:                                                                                                       \
    void set_##field_name(data_type field_name)                                                               \
    {                                                                                                         \
        field_name##_ = field_name;                                                                           \
        *(reinterpret_cast<data_type *>(dataPtr_ + field_name##_offset_)) = field_name;                       \
    }                                                                                                        \
    data_type get_##field_name()                                                                              \
    {                                                                                                         \
        return field_name##_;                                                                                 \
    }                                                                                                         \
                                                                                                              \
private:                                                                                                      \
    uint8_t get_##field_name##_reserve_buf() const                                                            \
    {                                                                                                         \
        return field_name##_reserve_buf_[0];                                                                  \
    }                                                                                                         \
    data_type field_name##_ = 0;                                                                              \
    size_t field_name##_offset_ = FieldHandler(#data_type, #field_name, sizeof(data_type), #field_name "PH"); \
    uint8_t field_name##_reserve_buf_[16] = {0}

#define TILING_DATA_FIELD_DEF_ARR(arr_type, arr_size, field_name)                                                      \
public:                                                                                                                \
    void set_##field_name(arr_type *field_name)                                                                        \
    {                                                                                                                  \
        field_name##_ = field_name;                                                                                    \
        auto offset = field_name##_offset_;                                                                            \
        if (dataPtr_ + offset == (uint8_t *)field_name) {                                                              \
            return;                                                                                                    \
        }                                                                                                              \
        const auto err_t = memcpy_s(dataPtr_ + offset, dataSize_ - offset, field_name, (arr_size) * sizeof(arr_type)); \
        if (err_t != EOK) {                                                                                            \
            GeLogError("tilingdata_base.h TILING_DATA_FIELD_DEF_ARR memcpy is failed !");                              \
        }                                                                                                              \
    }                                                                                                                  \
    arr_type *get_##field_name()                                                                                       \
    {                                                                                                                  \
        return (arr_type *)(dataPtr_ + field_name##_offset_);                                                          \
    }                                                                                                                  \
                                                                                                                       \
private:                                                                                                               \
    arr_type *field_name##_ = nullptr;                                                                                 \
    size_t field_name##_offset_ = FieldHandler(#arr_type, #field_name, sizeof(arr_type), arr_size, #field_name "PH")

#define TILING_DATA_FIELD_DEF_STRUCT(struct_type, field_name)          \
public:                                                                \
    struct_type field_name{nullptr};                                   \
                                                                       \
private:                                                               \
    size_t field_name##_offset_ = FieldHandler("struct",               \
        #field_name,                                                   \
        #struct_type,                                                  \
        StructSizeInfoBase::GetInstance().GetStructSize(#struct_type), \
        reinterpret_cast<void *>(&field_name),                                           \
        #field_name "PH")

#define END_TILING_DATA_DEF \
    }

#define REGISTER_TILING_DATA_CLASS(opType, className)                                      \
    class opType##className##Helper {                                                      \
    public:                                                                                \
        opType##className##Helper()                                                        \
        {                                                                                  \
            CTilingDataClassFactory::GetInstance().RegisterTilingData(                     \
                #opType, #className, opType##className##Helper::CreateTilingDataInstance); \
        }                                                                                  \
        static std::shared_ptr<TilingDef> CreateTilingDataInstance()                       \
        {                                                                                  \
            return std::make_shared<className>();                                          \
        }                                                                                  \
    };                                                                                     \
    static className g_##opType##className##init;                                          \
    static opType##className##Helper g_tilingdata_##opType##className##helper;

#endif  // GE_TILINGDATA_BASE_H